'''
拖入视频，把视频转换成图片，存在当前目录下
'''

import tkinterdnd2
from tkinter import ttk, messagebox
import cv2
import os
from PIL import Image
import threading


class FrameExtractor(tkinterdnd2.Tk):
    def __init__(self):
        super().__init__()

        self.title("视频帧提取工具")
        self.geometry("800x600")

        # 创建主框架
        self.main_frame = ttk.Frame(self)
        self.main_frame.pack(expand=True, fill='both', padx=10, pady=10)

        # 创建裁剪设置区域
        self.settings_frame = ttk.LabelFrame(self.main_frame, text="裁剪设置")
        self.settings_frame.pack(fill='x', padx=5, pady=5)

        # 创建设置输入框
        self.create_settings_inputs()

        # 创建拖放区域
        self.drop_frame = ttk.LabelFrame(self.main_frame, text="将视频文件拖放到这里")
        self.drop_frame.pack(expand=True, fill='both', padx=5, pady=5)

        # 创建状态标签
        self.status_label = ttk.Label(self.drop_frame, text="等待视频文件...", anchor='center')
        self.status_label.pack(expand=True)

        # 创建进度条
        self.progress = ttk.Progressbar(self.main_frame, mode='determinate')
        self.progress.pack(fill='x', padx=5, pady=5)

        # 配置拖放区域
        self.drop_frame.drop_target_register(tkinterdnd2.DND_FILES)
        self.drop_frame.dnd_bind('<<Drop>>', self.handle_drop)
        self.drop_frame.dnd_bind('<<DragEnter>>', self.handle_drag_enter)
        self.drop_frame.dnd_bind('<<DragLeave>>', self.handle_drag_leave)

    def create_settings_inputs(self):
        # 创建输入框容器
        input_frame = ttk.Frame(self.settings_frame)
        input_frame.pack(fill='x', padx=5, pady=5)

        # X坐标输入
        ttk.Label(input_frame, text="X:").pack(side='left', padx=(0, 2))
        self.x_var = tk.StringVar(value="0")
        ttk.Entry(input_frame, textvariable=self.x_var, width=8).pack(side='left', padx=(0, 10))

        # Y坐标输入
        ttk.Label(input_frame, text="Y:").pack(side='left', padx=(0, 2))
        self.y_var = tk.StringVar(value="0")
        ttk.Entry(input_frame, textvariable=self.y_var, width=8).pack(side='left', padx=(0, 10))

        # 宽度输入
        ttk.Label(input_frame, text="宽度:").pack(side='left', padx=(0, 2))
        self.width_var = tk.StringVar(value="1440")
        ttk.Entry(input_frame, textvariable=self.width_var, width=8).pack(side='left', padx=(0, 10))

        # 高度输入
        ttk.Label(input_frame, text="高度:").pack(side='left', padx=(0, 2))
        self.height_var = tk.StringVar(value="720")
        ttk.Entry(input_frame, textvariable=self.height_var, width=8).pack(side='left', padx=(0, 10))

        # 创建格式选择下拉框
        ttk.Label(input_frame, text="格式:").pack(side='left', padx=(0, 2))
        self.format_var = tk.StringVar(value="jpg")
        format_combo = ttk.Combobox(input_frame, textvariable=self.format_var, values=["jpg", "png"], width=5)
        format_combo.pack(side='left', padx=(0, 10))

        # 创建质量设置（仅用于JPG）
        ttk.Label(input_frame, text="质量:").pack(side='left', padx=(0, 2))
        self.quality_var = tk.StringVar(value="95")
        ttk.Entry(input_frame, textvariable=self.quality_var, width=4).pack(side='left')

    def get_crop_settings(self):
        try:
            settings = {
                'x': int(self.x_var.get()),
                'y': int(self.y_var.get()),
                'width': int(self.width_var.get()),
                'height': int(self.height_var.get()),
                'format': self.format_var.get(),
                'quality': int(self.quality_var.get())
            }
            return settings
        except ValueError:
            messagebox.showerror("错误", "请输入有效的数值")
            return None

    def handle_drag_enter(self, event):
        self.drop_frame.configure(relief='solid')

    def handle_drag_leave(self, event):
        self.drop_frame.configure(relief='groove')

    def handle_drop(self, event):
        file_path = event.data
        
        # Windows系统可能会在文件路径两端加上大括号，需要去除
        if file_path.startswith('{') and file_path.endswith('}'):
            file_path = file_path[1:-1]
        
        # 检查文件是否存在
        if not os.path.exists(file_path):
            messagebox.showerror("错误", "文件不存在")
            return
            
        # 检查文件扩展名
        _, ext = os.path.splitext(file_path)
        if ext.lower() not in ['.mp4', '.avi', '.mov', '.mkv']:
            messagebox.showerror("错误", "不支持的文件格式")
            return

        # 获取裁剪设置
        settings = self.get_crop_settings()
        if settings is None:
            return

        # 在新线程中处理视频
        thread = threading.Thread(target=self.process_video, args=(file_path, settings))
        thread.daemon = True
        thread.start()

    def process_video(self, video_path, settings):
        try:
            self.status_label.config(text="正在处理视频...")
            self.progress['value'] = 0
            
            # 打开视频文件
            cap = cv2.VideoCapture(video_path)
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            
            # 创建输出目录
            video_dir = os.path.dirname(video_path)
            output_dir = os.path.join(video_dir, 'pictures')
            os.makedirs(output_dir, exist_ok=True)
            
            frame_count = 0
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                    
                # 裁剪图片
                x, y = settings['x'], settings['y']
                w, h = settings['width'], settings['height']
                
                # 检查裁剪区域是否超出原图范围
                orig_h, orig_w = frame.shape[:2]
                if x + w > orig_w or y + h > orig_h:
                    raise ValueError(f"裁剪区域超出原图范围 (原图尺寸: {orig_w}x{orig_h})")
                
                frame = frame[y:y+h, x:x+w]
                
                # 转换颜色空间
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                
                # 保存图片
                img = Image.fromarray(frame)
                # output_path = os.path.join(output_dir, f'frame_{frame_count:06d}.{settings["format"]}')
                output_path = os.path.join(output_dir, f'{frame_count+1}.{settings["format"]}')
                img.save(output_path, quality=settings['quality'])
                
                # 更新进度
                frame_count += 1
                progress = (frame_count / total_frames) * 100
                self.progress['value'] = progress
                self.status_label.config(text=f"已处理: {frame_count}/{total_frames} 帧")
                self.update_idletasks()
            
            cap.release()
            self.status_label.config(text=f"处理完成! 共提取 {frame_count} 帧")
            # messagebox.showinfo("完成", f"视频处理完成!\n输出目录: {output_dir}")
            
        except Exception as e:
            messagebox.showerror("错误", f"处理视频时出错: {str(e)}")
            self.status_label.config(text="处理失败")
        
        self.progress['value'] = 0

if __name__ == "__main__":
    import tkinter as tk
    app = FrameExtractor()
    app.mainloop()